#!/bin/bash

# AIQuant 开发环境部署脚本
# 用途: 快速搭建和管理开发环境
# 作者: AIQuant Team
# 版本: 1.0.0

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 配置变量
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
DEPLOY_ENV="development"
COMPOSE_FILE="docker-compose.yml"
DEV_DATA_DIR="$PROJECT_ROOT/dev_data"
LOG_DIR="$PROJECT_ROOT/logs"

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 错误处理
error_exit() {
    log_error "$1"
    exit 1
}

# 显示帮助信息
show_help() {
    cat << EOF
AIQuant 开发环境部署脚本

用法: $0 [COMMAND] [OPTIONS]

命令:
    start           启动开发环境 (默认)
    stop            停止开发环境
    restart         重启开发环境
    reset           重置开发环境 (删除所有数据)
    logs            查看服务日志
    status          查看服务状态
    shell           进入API容器shell
    db-shell        进入数据库shell
    test            运行测试
    build           构建镜像
    clean           清理资源

选项:
    -h, --help      显示帮助信息
    -v, --verbose   详细输出
    -d, --detach    后台运行
    -f, --force     强制执行
    --no-cache      构建时不使用缓存
    --pull          拉取最新基础镜像

示例:
    $0 start                    # 启动开发环境
    $0 start --verbose          # 启动并显示详细日志
    $0 reset --force            # 强制重置环境
    $0 logs api                 # 查看API服务日志
    $0 shell                    # 进入API容器

EOF
}

# 解析命令行参数
parse_args() {
    COMMAND="start"
    SERVICE=""
    
    while [[ $# -gt 0 ]]; do
        case $1 in
            start|stop|restart|reset|logs|status|shell|db-shell|test|build|clean)
                COMMAND="$1"
                shift
                ;;
            -h|--help)
                show_help
                exit 0
                ;;
            -v|--verbose)
                VERBOSE=true
                shift
                ;;
            -d|--detach)
                DETACH=true
                shift
                ;;
            -f|--force)
                FORCE=true
                shift
                ;;
            --no-cache)
                NO_CACHE=true
                shift
                ;;
            --pull)
                PULL=true
                shift
                ;;
            -*)
                error_exit "未知选项: $1"
                ;;
            *)
                if [[ -z "$SERVICE" ]]; then
                    SERVICE="$1"
                fi
                shift
                ;;
        esac
    done
}

# 检查依赖
check_dependencies() {
    log_info "检查依赖..."
    
    # 检查必需的命令
    local required_commands=("docker" "docker-compose" "curl" "jq")
    for cmd in "${required_commands[@]}"; do
        if ! command -v "$cmd" &> /dev/null; then
            error_exit "缺少必需命令: $cmd"
        fi
    done
    
    # 检查Docker服务
    if ! docker info &> /dev/null; then
        error_exit "Docker服务未运行"
    fi
    
    # 检查Docker Compose版本
    local compose_version=$(docker-compose --version | grep -oE '[0-9]+\.[0-9]+\.[0-9]+' | head -1)
    local min_version="1.27.0"
    if [[ "$(printf '%s\n' "$min_version" "$compose_version" | sort -V | head -n1)" != "$min_version" ]]; then
        log_warning "Docker Compose版本过低，建议升级到$min_version或更高版本"
    fi
    
    log_success "依赖检查完成"
}

# 初始化开发环境
init_dev_environment() {
    log_info "初始化开发环境..."
    
    cd "$PROJECT_ROOT"
    
    # 创建必要的目录
    mkdir -p "$DEV_DATA_DIR"/{postgres,redis,timescaledb,uploads,logs}
    mkdir -p "$LOG_DIR"
    mkdir -p data/{market_data,backtest_results,models}
    
    # 复制环境配置文件
    if [[ ! -f ".env" ]]; then
        if [[ -f ".env.example" ]]; then
            cp .env.example .env
            log_info "已创建.env文件，请根据需要修改配置"
        else
            log_warning ".env.example文件不存在，请手动创建.env文件"
        fi
    fi
    
    # 设置文件权限
    chmod -R 755 "$DEV_DATA_DIR"
    chmod -R 755 "$LOG_DIR"
    
    # 创建开发用的数据库初始化脚本
    cat > "$DEV_DATA_DIR/init_dev_db.sql" << 'EOF'
-- 开发环境数据库初始化
CREATE DATABASE IF NOT EXISTS aiquant_dev;
CREATE DATABASE IF NOT EXISTS aiquant_test;

-- 创建开发用户
CREATE USER IF NOT EXISTS 'dev_user'@'%' IDENTIFIED BY 'dev_password';
GRANT ALL PRIVILEGES ON aiquant_dev.* TO 'dev_user'@'%';
GRANT ALL PRIVILEGES ON aiquant_test.* TO 'dev_user'@'%';

-- 创建测试数据
USE aiquant_dev;

-- 插入示例用户
INSERT IGNORE INTO users (id, username, email, created_at) VALUES 
(1, 'dev_user', 'dev@aiquant.com', NOW()),
(2, 'test_user', 'test@aiquant.com', NOW());

-- 插入示例策略
INSERT IGNORE INTO strategies (id, name, description, user_id, created_at) VALUES 
(1, '示例策略1', '这是一个示例量化策略', 1, NOW()),
(2, '示例策略2', '这是另一个示例策略', 1, NOW());

FLUSH PRIVILEGES;
EOF
    
    log_success "开发环境初始化完成"
}

# 构建镜像
build_images() {
    log_info "构建Docker镜像..."
    
    cd "$PROJECT_ROOT"
    
    local build_args=""
    if [[ "$NO_CACHE" == "true" ]]; then
        build_args="--no-cache"
    fi
    
    if [[ "$PULL" == "true" ]]; then
        build_args="$build_args --pull"
    fi
    
    # 构建主应用镜像
    log_info "构建主应用镜像..."
    docker-compose build $build_args api
    
    # 构建Web界面镜像
    log_info "构建Web界面镜像..."
    docker-compose build $build_args web
    
    # 构建Worker镜像
    log_info "构建Worker镜像..."
    docker-compose build $build_args worker
    
    log_success "镜像构建完成"
}

# 启动服务
start_services() {
    log_info "启动开发环境..."
    
    cd "$PROJECT_ROOT"
    
    # 检查配置文件
    if [[ ! -f "$COMPOSE_FILE" ]]; then
        error_exit "Docker Compose文件不存在: $COMPOSE_FILE"
    fi
    
    # 验证配置
    if ! docker-compose config &> /dev/null; then
        error_exit "Docker Compose配置有误"
    fi
    
    # 启动基础服务（数据库、缓存等）
    log_info "启动基础服务..."
    docker-compose up -d postgres redis timescaledb
    
    # 等待数据库就绪
    log_info "等待数据库就绪..."
    local attempt=1
    while [[ $attempt -le 30 ]]; do
        if docker-compose exec -T postgres pg_isready -U aiquant; then
            break
        fi
        log_info "等待数据库启动... ($attempt/30)"
        sleep 2
        ((attempt++))
    done
    
    if [[ $attempt -gt 30 ]]; then
        error_exit "数据库启动超时"
    fi
    
    # 运行数据库迁移
    log_info "执行数据库迁移..."
    docker-compose run --rm api python scripts/migrate_db.py || log_warning "数据库迁移失败"
    
    # 启动应用服务
    log_info "启动应用服务..."
    if [[ "$DETACH" == "true" ]]; then
        docker-compose up -d
    else
        if [[ "$VERBOSE" == "true" ]]; then
            docker-compose up
        else
            docker-compose up -d
            log_info "服务已在后台启动，使用 '$0 logs' 查看日志"
        fi
    fi
    
    # 等待服务就绪
    log_info "等待服务就绪..."
    sleep 10
    
    # 健康检查
    if check_services_health; then
        log_success "开发环境启动完成！"
        show_service_urls
    else
        log_warning "部分服务可能未正常启动，请检查日志"
    fi
}

# 停止服务
stop_services() {
    log_info "停止开发环境..."
    
    cd "$PROJECT_ROOT"
    docker-compose down
    
    log_success "开发环境已停止"
}

# 重启服务
restart_services() {
    log_info "重启开发环境..."
    
    stop_services
    sleep 5
    start_services
}

# 重置环境
reset_environment() {
    if [[ "$FORCE" != "true" ]]; then
        log_warning "此操作将删除所有开发数据！"
        read -p "确认重置环境？(y/N): " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            log_info "重置已取消"
            return 0
        fi
    fi
    
    log_info "重置开发环境..."
    
    cd "$PROJECT_ROOT"
    
    # 停止并删除容器
    docker-compose down -v --remove-orphans
    
    # 删除开发数据
    if [[ -d "$DEV_DATA_DIR" ]]; then
        rm -rf "$DEV_DATA_DIR"
        log_info "已删除开发数据目录"
    fi
    
    # 删除日志
    if [[ -d "$LOG_DIR" ]]; then
        rm -rf "$LOG_DIR"
        log_info "已删除日志目录"
    fi
    
    # 清理Docker资源
    docker system prune -f
    
    # 重新初始化
    init_dev_environment
    
    log_success "环境重置完成"
}

# 查看日志
show_logs() {
    cd "$PROJECT_ROOT"
    
    if [[ -n "$SERVICE" ]]; then
        log_info "查看 $SERVICE 服务日志..."
        docker-compose logs -f "$SERVICE"
    else
        log_info "查看所有服务日志..."
        docker-compose logs -f
    fi
}

# 查看服务状态
show_status() {
    log_info "服务状态："
    
    cd "$PROJECT_ROOT"
    docker-compose ps
    
    echo
    log_info "资源使用情况："
    docker stats --no-stream --format "table {{.Container}}\t{{.CPUPerc}}\t{{.MemUsage}}\t{{.NetIO}}\t{{.BlockIO}}"
}

# 进入容器shell
enter_shell() {
    cd "$PROJECT_ROOT"
    
    local container="api"
    if [[ -n "$SERVICE" ]]; then
        container="$SERVICE"
    fi
    
    log_info "进入 $container 容器..."
    docker-compose exec "$container" /bin/bash
}

# 进入数据库shell
enter_db_shell() {
    cd "$PROJECT_ROOT"
    
    log_info "进入数据库shell..."
    docker-compose exec postgres psql -U aiquant -d aiquant
}

# 运行测试
run_tests() {
    log_info "运行测试..."
    
    cd "$PROJECT_ROOT"
    
    # 确保测试数据库存在
    docker-compose exec -T postgres createdb -U aiquant aiquant_test 2>/dev/null || true
    
    # 运行单元测试
    log_info "运行单元测试..."
    docker-compose run --rm -e DATABASE_URL="postgresql://aiquant:password@postgres:5432/aiquant_test" \
        api python -m pytest tests/unit/ -v
    
    # 运行集成测试
    log_info "运行集成测试..."
    docker-compose run --rm -e DATABASE_URL="postgresql://aiquant:password@postgres:5432/aiquant_test" \
        api python -m pytest tests/integration/ -v
    
    # 运行API测试
    log_info "运行API测试..."
    docker-compose run --rm api python -m pytest tests/api/ -v
    
    log_success "测试完成"
}

# 清理资源
clean_resources() {
    log_info "清理开发环境资源..."
    
    cd "$PROJECT_ROOT"
    
    # 停止容器
    docker-compose down --remove-orphans
    
    # 清理未使用的镜像
    docker image prune -f
    
    # 清理未使用的卷
    docker volume prune -f
    
    # 清理未使用的网络
    docker network prune -f
    
    # 清理构建缓存
    docker builder prune -f
    
    log_success "资源清理完成"
}

# 检查服务健康状态
check_services_health() {
    local all_healthy=true
    
    # 检查API服务
    if curl -f -s http://localhost:8000/health > /dev/null; then
        log_success "API服务正常"
    else
        log_warning "API服务异常"
        all_healthy=false
    fi
    
    # 检查Web服务
    if curl -f -s http://localhost:8501 > /dev/null; then
        log_success "Web服务正常"
    else
        log_warning "Web服务异常"
        all_healthy=false
    fi
    
    # 检查数据库
    cd "$PROJECT_ROOT"
    if docker-compose exec -T postgres pg_isready -U aiquant; then
        log_success "数据库连接正常"
    else
        log_warning "数据库连接异常"
        all_healthy=false
    fi
    
    # 检查Redis
    if docker-compose exec -T redis redis-cli ping | grep -q PONG; then
        log_success "Redis连接正常"
    else
        log_warning "Redis连接异常"
        all_healthy=false
    fi
    
    return $all_healthy
}

# 显示服务URL
show_service_urls() {
    echo
    log_info "服务访问地址："
    echo "  API服务:     http://localhost:8000"
    echo "  API文档:     http://localhost:8000/docs"
    echo "  Web界面:     http://localhost:8501"
    echo "  监控面板:    http://localhost:3000 (如果启用)"
    echo "  数据库:      localhost:5432"
    echo "  Redis:       localhost:6379"
    echo
    log_info "常用命令："
    echo "  查看日志:    $0 logs [service]"
    echo "  查看状态:    $0 status"
    echo "  进入容器:    $0 shell [service]"
    echo "  运行测试:    $0 test"
    echo "  停止服务:    $0 stop"
}

# 安装开发工具
install_dev_tools() {
    log_info "安装开发工具..."
    
    # 安装pre-commit hooks
    if command -v pre-commit &> /dev/null; then
        pre-commit install
        log_success "Pre-commit hooks已安装"
    else
        log_warning "Pre-commit未安装，跳过hooks安装"
    fi
    
    # 创建开发用的别名脚本
    cat > "$PROJECT_ROOT/dev_aliases.sh" << 'EOF'
#!/bin/bash
# 开发环境别名

alias aiq-start="./scripts/deploy/development_deploy.sh start"
alias aiq-stop="./scripts/deploy/development_deploy.sh stop"
alias aiq-restart="./scripts/deploy/development_deploy.sh restart"
alias aiq-logs="./scripts/deploy/development_deploy.sh logs"
alias aiq-status="./scripts/deploy/development_deploy.sh status"
alias aiq-shell="./scripts/deploy/development_deploy.sh shell"
alias aiq-test="./scripts/deploy/development_deploy.sh test"
alias aiq-reset="./scripts/deploy/development_deploy.sh reset"

echo "AIQuant开发环境别名已加载！"
echo "使用 aiq-start, aiq-stop, aiq-logs 等命令快速操作"
EOF
    
    chmod +x "$PROJECT_ROOT/dev_aliases.sh"
    log_info "开发别名脚本已创建: dev_aliases.sh"
    log_info "运行 'source dev_aliases.sh' 加载别名"
}

# 主函数
main() {
    log_info "AIQuant开发环境管理"
    log_info "命令: $COMMAND"
    log_info "时间: $(date)"
    
    # 解析命令行参数
    parse_args "$@"
    
    # 检查依赖
    check_dependencies
    
    # 初始化环境（如果需要）
    if [[ ! -d "$DEV_DATA_DIR" ]]; then
        init_dev_environment
    fi
    
    # 执行命令
    case "$COMMAND" in
        start)
            start_services
            ;;
        stop)
            stop_services
            ;;
        restart)
            restart_services
            ;;
        reset)
            reset_environment
            ;;
        logs)
            show_logs
            ;;
        status)
            show_status
            ;;
        shell)
            enter_shell
            ;;
        db-shell)
            enter_db_shell
            ;;
        test)
            run_tests
            ;;
        build)
            build_images
            ;;
        clean)
            clean_resources
            ;;
        *)
            error_exit "未知命令: $COMMAND"
            ;;
    esac
}

# 错误处理
set -E
trap 'log_error "执行过程中发生错误，行号: $LINENO"; exit 1' ERR

# 执行主函数
main "$@"